{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cv2 as cv\n",
    "import numpy as np\n",
    "import time\n",
    "import math\n",
    "\n",
    "import ipywidgets as widgets\n",
    "from IPython.display import display"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "straightLaneVideo = cv.VideoCapture('data/StraightLane.mp4')\n",
    "turnRightVideo = cv.VideoCapture('data/TurnRightLane.mp4')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "maxWidth = 640\n",
    "maxHeight = 480\n",
    "windowsObj = {}\n",
    "def imshow(name, img):\n",
    "    if(not name in windowsObj):\n",
    "        windowsObj[name] = widgets.Image(format='jpg', height=maxHeight, width=maxWidth)\n",
    "        display(windowsObj[name])\n",
    "    windowsObj[name].value = cv.imencode('.jpg', img)[1].tobytes()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "currentTime = time.time()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def colorDetector(image, lowerThreshold, upperThreshold):\n",
    "    hsvImg = cv.cvtColor(image, cv.COLOR_RGB2HSV)\n",
    "    maskImg = cv.inRange(hsvImg, lowerThreshold, upperThreshold)\n",
    "    \n",
    "    kernel = np.ones((5, 5), np.uint8)\n",
    "    maskImg = cv.morphologyEx(maskImg, cv.MORPH_OPEN, kernel)\n",
    "    \n",
    "    roiPoints = np.array([[0, maxHeight],\n",
    "                         [maxWidth, maxHeight],\n",
    "                         [maxWidth, maxHeight/2],\n",
    "                         [0, maxHeight/2]], np.int32)\n",
    "    roiMask = np.zeros((maxHeight, maxWidth), np.uint8)\n",
    "    cv.fillPoly(roiMask, [roiPoints], (255))\n",
    "    \n",
    "    maskImg = cv.bitwise_and(maskImg, roiMask)\n",
    "    \n",
    "    maskImg = cv.medianBlur(maskImg, 5)\n",
    "    return maskImg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fitLine(maskImg):\n",
    "    lineEdgesImg = cv.Canny(maskImg, 20, 60)\n",
    "    \n",
    "    lineContours, _ = cv.findContours(lineEdgesImg, cv.RETR_TREE, cv.CHAIN_APPROX_SIMPLE)\n",
    "    W = 0\n",
    "    b = -999\n",
    "    if(len(lineContours) > 0):\n",
    "        mID = 0\n",
    "        mA = 0\n",
    "        for i in range(len(lineContours)):\n",
    "            tmp = cv.contourArea(lineContours[i])\n",
    "            if(tmp > mA):\n",
    "                mA = tmp\n",
    "                mID = i\n",
    "        \n",
    "        M = cv.moments(lineContours[mID])\n",
    "        if(M['m00'] != 0):\n",
    "            cX = int(M['m10'] / M['m00'])\n",
    "            cY = int(M['m01'] / M['m00'])\n",
    "            \n",
    "            rect = cv.minAreaRect(lineContours[mID])\n",
    "            box = cv.boxPoints(rect)\n",
    "            \n",
    "            def calDist(p1, p2):\n",
    "                return math.sqrt((p1[0] - p2[0]) * (p1[0] - p2[0]) + (p1[1]-p2[1]) * (p1[1] - p2[1]))\n",
    "            \n",
    "            if(calDist(box[0], box[1]) > calDist(box[1], box[2]) and box[1][0] - box[0][0] != 0):\n",
    "                W = (box[1][1] - box[0][1]) / (box[1][0] - box[0][0])\n",
    "            elif(box[2][0] - box[1][0] != 0):\n",
    "                W = (box[2][1] - box[1][1]) / (box[2][0] - box[1][0])\n",
    "            \n",
    "            b = cY - W * cX\n",
    "    return W, b  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = np.empty([0, 2])\n",
    "labels = np.empty([0])\n",
    "cnt1 = 0\n",
    "cnt2 = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "db546da9f70f451a83c51afee62beb81",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Image(value=b'', format='jpg', height='480', width='640')"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "40d95aadcee04ea790cc8e6e3ed1d936",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Image(value=b'', format='jpg', height='480', width='640')"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "bbd18fdf0bc340839d83961e67f6da06",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Image(value=b'', format='jpg', height='480', width='640')"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "video alreay ended.\n"
     ]
    }
   ],
   "source": [
    "while 1:\n",
    "    if(time.time() - currentTime > 0.033):\n",
    "        currentTie = time.time()\n",
    "        ret,frame = straightLaneVideo.read()\n",
    "        \n",
    "        if(ret):\n",
    "            imshow('straightLaneFrame', frame)\n",
    "            rgbImg = cv.cvtColor(frame, cv.COLOR_BGR2RGB)\n",
    "            \n",
    "            lowerThreshold = np.array([11, 80, 90])\n",
    "            upperThreshold = np.array([35, 255, 2555])\n",
    "            \n",
    "            maskImg = colorDetector(rgbImg, lowerThreshold, upperThreshold)\n",
    "            imshow('straightLaneMask', maskImg)\n",
    "            \n",
    "            W, b = fitLine(maskImg)\n",
    "            \n",
    "            if(W != 0 and b != -999):\n",
    "                y1 = 0\n",
    "                x1 = int((y1 - b) / W)\n",
    "                y2 = maxHeight\n",
    "                x2 = int((y2-b) / W)\n",
    "                 \n",
    "                cv.line(frame, (x1, y1), (x2, y2), (255), 3)\n",
    "                imshow('straightLane', frame)\n",
    "                \n",
    "                cnt1 += 1\n",
    "                data = np.concatenate((data, np.array([[W, b]])), axis = 0)\n",
    "        else:\n",
    "            print('video alreay ended.')\n",
    "            break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "if(cnt1 > 0):\n",
    "    label = np.full(cnt1, 0)\n",
    "    labels = np.append(labels, label)\n",
    "    \n",
    "straightLaneVideo.release()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "406b22d5cc634809bcb5986ab56f2cb8",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Image(value=b'', format='jpg', height='480', width='640')"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "36186fa907a84702bff2a212046240c5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Image(value=b'', format='jpg', height='480', width='640')"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a0b85fde49464786976a2145e2972671",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Image(value=b'', format='jpg', height='480', width='640')"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "while 1:\n",
    "    if(time.time() - currentTime > 0.033):\n",
    "        currentTime = time.time()\n",
    "\n",
    "        ret, frame = turnRightVideo.read()\n",
    "\n",
    "        if(ret):\n",
    "            imshow('turnRightLaneFrame', frame)\n",
    "\n",
    "            rgbImg = cv.cvtColor(frame, cv.COLOR_BGR2RGB)\n",
    "\n",
    "            lowerThreshold = np.array([11, 80, 90])\n",
    "            upperThreshold = np.array([35, 255, 255])\n",
    "\n",
    "            maskImg = colorDetector(rgbImg, lowerThreshold, upperThreshold)\n",
    "            imshow('turnRightLaneMask', maskImg)\n",
    "\n",
    "            W,b = fitLine(maskImg)\n",
    "\n",
    "            if(W != 0 and b != -999):\n",
    "                y1 = 0\n",
    "                x1 = int((y1 - b) / W)\n",
    "\n",
    "                y2 = maxHeight\n",
    "                x2 = int((y2 - b) / W)\n",
    "\n",
    "                cv.line(frame, (x1,y1), (x2,y2), (255), 3)\n",
    "                imshow('turnRightLane', frame)\n",
    "\n",
    "                cnt2 += 1\n",
    "                data = np.concatenate((data, np.array([[W, b]])), axis=0)\n",
    "\n",
    "        else:\n",
    "            print('视频已播放完毕')\n",
    "            break\n",
    "\n",
    "if(cnt2 > 0):\n",
    "    label = np.full(cnt2, 1)\n",
    "    labels = np.append(labels, label)\n",
    "\n",
    "turnRightVideo.release()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(data.shape)\n",
    "print(labels.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.save('data/data.npy', data)\n",
    "np.save('data/labels.npy', labels)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
